package com.etl.component.file;

import cn.hutool.core.util.StrUtil;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.operators.DataSink;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.operators.MapOperator;
import org.apache.flink.api.java.operators.UnionOperator;
import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.table.sinks.CsvTableSink;
import org.apache.flink.table.utils.TableConnectorUtils;
import org.apache.flink.types.Row;

/**
 * 描述：
 * flink这个版本的csv解析不能识别特殊字符串和表头 这里做一个统一处理
 * @author xianggj
 * @Date 2021/11/8 16:17
 **/
public class MyCsvTableSink extends CsvTableSink {

    /**
     * 特殊转义字符串
     */
    private Character quotedStringParsing = '\"';

    /**
     * 特殊转义字符串
     */
    private boolean enableHeader;

    /**
     * 继承不过来的私有属性
     */
    private FileSystem.WriteMode writeMode;
    private String path;
    private int numFiles;
    private String fieldDelim;

    public MyCsvTableSink(String path, String fieldDelim, int numFiles, FileSystem.WriteMode writeMode) {
        super(path, fieldDelim, numFiles, writeMode);
        this.path = path;
        this.numFiles = numFiles;
        this.writeMode = writeMode;
        this.fieldDelim = fieldDelim == null ? "," : fieldDelim;
    }

    public void setQuotedStringParsing(Character quotedStringParsing) {
        this.quotedStringParsing = quotedStringParsing;
    }

    public void setEnableHeader(boolean enableHeader) {
        this.enableHeader = enableHeader;
    }

    /**
     * 暂时先复写一个
     * @param dataSet
     * @return
     */
    @Override
    public DataSink<?> consumeDataSet(DataSet<Row> dataSet) {
        MapOperator<Row, String> csvRows = dataSet.map(
                new MyCsvTableSink.CsvFormatter(this.fieldDelim, this.quotedStringParsing));
        DataSink sink;
        if (enableHeader) {
            RowTypeInfo rowTypeInfo = (RowTypeInfo)dataSet.getType();
            String[] fieldNames = rowTypeInfo.getFieldNames();
            String join = StrUtil.join(this.fieldDelim, fieldNames);
            DataSource<String> header = dataSet.getExecutionEnvironment().fromElements(join);
            UnionOperator<String> union = header.union(csvRows);
            if (this.writeMode != null) {
                sink = union.writeAsText(this.path, this.writeMode);
            } else {
                sink = union.writeAsText(this.path);
            }
            if (this.numFiles > 0) {
                sink.setParallelism(this.numFiles);
            }
        } else {
            if (this.writeMode != null) {
                sink = csvRows.writeAsText(this.path, this.writeMode);
            } else {
                sink = csvRows.writeAsText(this.path);
            }
            if (this.numFiles > 0) {
                csvRows.setParallelism(this.numFiles);
                sink.setParallelism(this.numFiles);
            }
        }


        return sink.name(TableConnectorUtils.generateRuntimeName(CsvTableSink.class, null));
    }

    public static class CsvFormatter implements MapFunction<Row, String> {
        private static final long serialVersionUID = 1L;
        private final String fieldDelim;
        private final Character quotedStringParsing;

        CsvFormatter(String fieldDelim, Character quotedStringParsing) {
            this.fieldDelim = fieldDelim;
            this.quotedStringParsing = quotedStringParsing;
        }

        @Override
        public String map(Row row) {
            StringBuilder builder = new StringBuilder();

            for(int i = 0; i < row.getArity(); ++i) {
                if (i > 0) {
                    builder.append(this.fieldDelim);
                }

                Object o;
                if ((o = row.getField(i)) != null) {
                    String s = String.valueOf(o);
                    if (s.contains(this.fieldDelim)){
                        builder.append(this.quotedStringParsing + s + this.quotedStringParsing);
                    } else {
                        builder.append(s);
                    }
                }
            }

            return builder.toString();
        }
    }
}
